import sys
sys.path.append("/data1/yxn/repast-city/")
import numpy as np
from core.env import Env
from core.vtype import ContinuousNormal, NamedCategorical
from envs.meituan_base import MeituanEnv


# 定义常量
POSITION_X = 'position_x'
POSITION_Y = 'position_y'
# RIDER_STATE = 'rider_state'
MOOD = 'mood'
REST_TIME = 'rest_time'
ORDER_POSITION_X = 'order_position_x'
ORDER_POSITION_Y = 'order_position_y'
ORDER_MONEY = 'order_money'
# ORDER_STATE = 'order_state'
PERFORMANCE = 'performance'

ACTION = 'action'

NEXT = {name: Env.name_next(name) for name in [POSITION_X, POSITION_Y, MOOD, REST_TIME, PERFORMANCE, ORDER_POSITION_X, ORDER_POSITION_Y]}

class MeituanRiderWrapper(Env):
    def __init__(self, args):
        super().__init__(args)
        self.__core = MeituanEnv()

    def define(self, args):
        _def = Env.Definition()
        self.__core = MeituanEnv()
        _def.state(POSITION_X,(ContinuousNormal(scale=None) ))
        _def.state(POSITION_Y,(ContinuousNormal(scale=None) ))
        # _def.state(RIDER_STATE, NamedCategoricalArray(list(NamedCategorical('idle', 'delivering','resting') for _ in range(self.__core.num_riders))))
        _def.state(MOOD, (ContinuousNormal(scale=None) ))
        _def.state(REST_TIME, (ContinuousNormal(scale=None) ))
        _def.state(PERFORMANCE,(ContinuousNormal(scale=None) ))
        _def.state(ORDER_MONEY,(ContinuousNormal(scale=None) ))
        _def.state(ORDER_POSITION_X,(ContinuousNormal(scale=None) ))
        _def.state(ORDER_POSITION_Y,(ContinuousNormal(scale=None) ))
        # _def.state(ORDER_STATE,NamedCategoricalArray(list(NamedCategorical('unprocessed', 'processing','completed') for _ in range(self.__core.num_orders))))
        _def.action(ACTION, NamedCategorical('accept_order', 'reject_order' ))
        # _def.reward('order completed', [NEXT[ORDER_STATE]], lambda x: np.sum(np.where((x == 2) | (x == 'completed'), 1.0, 0.0)))
        _def.reward('performance', [NEXT[PERFORMANCE]], lambda x:  np.mean(x / 100.0)) # 绩效标准化一下


        return _def
        
    @classmethod
    def init_parser(cls, parser):
        parser.add_argument("--render", action="store_true", default=False, help="render the environment in a window.")
    
    def launch(self):
        self.__core = MeituanEnv()

    def reset(self):
        super().reset()
        return self.__core.reset()
    
    def init_episode(self, *args, **kwargs):
        obs = self.__core.reset()
        return self._state_variables(obs)
    
    def transit(self, action):
        a = action[ACTION] #!!
        obs, rewards, done, info = self.__core.step(a)
        
        # 包装返回的状态和信息
        s = self._state_variables(obs)
        tran = {self.name_next(k): s[k] for k in s.keys()}
        
        return tran, info

    def _state_variables(self, obs):

    
        return {
            POSITION_X: obs[0],  # 位置x
            POSITION_Y: obs[1], # 位置y
            MOOD: obs[2],  # 骑手心情
            REST_TIME: obs[3],  # 骑手休息时间
            PERFORMANCE: obs[4] , # 骑手绩效
            ORDER_MONEY: obs[5],  # 订单金额
            ORDER_POSITION_X:obs[6], # 订单位置x
            ORDER_POSITION_Y:obs[7] #订单位置y
        }
    
    def terminated(self, transition) -> bool:
        return False  # 可以根据具体需求定义终止条件

    def random_action(self):
        return {ACTION: self.__core.action_space.sample()}
    
    def action_mask(self):
        return self.__core.action_mask()

    def __str__(self):
        return "Meituan"
# test
# env = MeituanRiderWrapper()
# obs = env.reset()
# print("Initial Observation:")
# print(obs)

# actions = [env.action_space.sample() for _ in range(env.num_riders)]
# obs, rewards, done, info = env.transit(actions)
# print("Next Observation:")
# print(obs)
# print("Rewards:")
# print(rewards)
